#
# This file is a part of the open-eBackup project.
# This Source Code Form is subject to the terms of the Mozilla Public License, v. 2.0.
# If a copy of the MPL was not distributed with this file, You can obtain one at
# http://mozilla.org/MPL/2.0/.
#
# Copyright (c) [2024] Huawei Technologies Co.,Ltd.
#
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
#

import struct
import base64

import hashlib
import random
import re
import binascii
from socket import gethostname

## 参考 [MS-NLMP].pdf https://winprotocoldoc.blob.core.windows.net/productionwindowsarchives/MS-NLMP/%5bMS-NLMP%5d.pdf

NTLMSSP_NEGOTIATE_UNICODE = 0x00000001
NTLM_NEGOTIATE_OEM = 0x00000002
NTLMSSP_REQUEST_TARGET = 0x00000004
NTLMSSP_NEGOTIATE_SIGN = 0x00000010
NTLMSSP_NEGOTIATE_SEAL = 0x00000020
NTLMSSP_NEGOTIATE_DATAGRAM = 0x00000040
NTLMSSP_NEGOTIATE_LM_KEY = 0x00000080
NTLMSSP_NEGOTIATE_NTLM = 0x00000200
NTLM_ANONYMOUS = 0x00000800
NTLMSSP_NEGOTIATE_OEM_DOMAIN_SUPPLIED = 0x00001000
NTLMSSP_NEGOTIATE_OEM_WORKSTATION_SUPPLIED = 0x00002000
NTLMSSP_NEGOTIATE_ALWAYS_SIGN = 0x00008000
NTLMSSP_TARGET_TYPE_DOMAIN = 0x00010000
NTLMSSP_TARGET_TYPE_SERVER = 0x00020000
NTLMSSP_NEGOTIATE_EXTENDED_SESSIONSECURITY = 0x00080000  # 如果设置，则表示使用NTLM v2，否则用NTLM v1
NTLMSSP_NEGOTIATE_IDENTIFY = 0x00100000
NTLMSSP_REQUEST_NON_NT_SESSION_KEY = 0x00400000
NTLMSSP_NEGOTIATE_TARGET_INFO = 0x00800000
NTLMSSP_NEGOTIATE_VERSION = 0x02000000
NTLMSSP_NEGOTIATE_128 = 0x20000000
NTLMSSP_NEGOTIATE_KEY_EXCH = 0x40000000
NTLMSSP_NEGOTIATE_56 = 0x80000000

NTLM_TYPE1_FLAGS = (NTLMSSP_NEGOTIATE_UNICODE |
                    NTLM_NEGOTIATE_OEM |
                    NTLMSSP_REQUEST_TARGET |
                    NTLMSSP_NEGOTIATE_NTLM |
                    NTLMSSP_NEGOTIATE_OEM_DOMAIN_SUPPLIED |
                    NTLMSSP_NEGOTIATE_OEM_WORKSTATION_SUPPLIED |
                    NTLMSSP_NEGOTIATE_ALWAYS_SIGN |
                    NTLMSSP_NEGOTIATE_EXTENDED_SESSIONSECURITY |
                    NTLMSSP_NEGOTIATE_VERSION |
                    NTLMSSP_NEGOTIATE_128 |
                    NTLMSSP_NEGOTIATE_56)
NTLM_TYPE2_FLAGS = (NTLMSSP_NEGOTIATE_UNICODE |
                    NTLMSSP_REQUEST_TARGET |
                    NTLMSSP_NEGOTIATE_NTLM |
                    NTLMSSP_NEGOTIATE_ALWAYS_SIGN |
                    NTLMSSP_NEGOTIATE_EXTENDED_SESSIONSECURITY |
                    NTLMSSP_NEGOTIATE_TARGET_INFO |
                    NTLMSSP_NEGOTIATE_VERSION |
                    NTLMSSP_NEGOTIATE_128 |
                    NTLMSSP_NEGOTIATE_56)

NTLMSSP_SIGNATURE = b'NTLMSSP\0'  # NTLM签名
NTLM_NEGOTIATE_MESSAGE_TYPE = struct.pack('<I', 1)  # NTLM协商报文类型
NTLM_AUTHENTICATE_MESSAGE_TYPE = struct.pack('<I', 3)  # NTLM认证报文类型
MAGIC_STR = b"KGS!@#$%"  # page 57 in [MS-NLMP]
NTLMSSP_VERSION_REVERSED_BIT = 0
NTLMSSP_REVISION_W2K3 = 15
NTLMSSP_VERSION_MAJOR = 5
NTLMSSP_VERSION_MINOR = 1
NTLMSSP_VERSION_PRODUCT_BUILD = 2600


def create_ntlm_negotiate_message(domain_name, type1_flags=NTLM_TYPE1_FLAGS):
    """
    生成ntlm negotiate 信息
    :param domain_name: 域
    :param type1_flags: 标志位
    :return: 信息
    """
    # 用于offset计算
    bytes_before_payload = 40
    payload_cur_offset = bytes_before_payload
    # 标志位
    flags = struct.pack('<I', type1_flags)
    # 生成workstation和domain_name
    workstation = bytes(gethostname().upper(), 'ascii')
    workstation_len = struct.pack('<H', len(workstation))
    workstation_max_len = struct.pack('<H', len(workstation))
    workstation_buffer_offset = struct.pack('<I', payload_cur_offset)
    payload_cur_offset += len(workstation)
    domain_name_bytes = bytes(domain_name, 'ascii')
    domain_name_len = struct.pack('<H', len(domain_name_bytes))
    domain_name_max_len = struct.pack('<H', len(domain_name_bytes))
    domain_name_buffer_offset = struct.pack('<I', payload_cur_offset)
    # 生成版本号
    version_str = get_version_str()
    # 生成结果
    base64_msg = base64.b64encode(NTLMSSP_SIGNATURE + NTLM_NEGOTIATE_MESSAGE_TYPE + flags +
                                  domain_name_len + domain_name_max_len + domain_name_buffer_offset +
                                  workstation_len + workstation_max_len + workstation_buffer_offset +
                                  version_str +
                                  workstation + domain_name_bytes)
    return base64_msg.decode()


def get_version_str():
    """
    生成版本字符串
    :return: 版本字符串
    """
    product_major_version = struct.pack('<B', NTLMSSP_VERSION_MAJOR)
    product_minor_version = struct.pack('<B', NTLMSSP_VERSION_MINOR)
    product_build = struct.pack('<H', NTLMSSP_VERSION_PRODUCT_BUILD)
    version_reserved1 = struct.pack('<B', NTLMSSP_VERSION_REVERSED_BIT)
    version_reserved2 = struct.pack('<B', NTLMSSP_VERSION_REVERSED_BIT)
    version_reserved3 = struct.pack('<B', NTLMSSP_VERSION_REVERSED_BIT)
    ntlm_revision_current = struct.pack('<B', NTLMSSP_REVISION_W2K3)
    version_str = product_major_version + product_minor_version + product_build + version_reserved1 + \
                  version_reserved2 + version_reserved3 + ntlm_revision_current
    return version_str


def parse_ntlm_challenge_message(msg2):
    """
    解析ntlm数据
    :param msg2: 源数据
    :return: 解析结果：(server_challenge, negotiate_flags)
    """
    msg2 = base64.decodebytes(bytes(msg2, 'ascii'))
    signature = msg2[0:8]
    if signature != NTLMSSP_SIGNATURE:
        raise Exception(f"signature not valid:{signature}")
    msg_type = struct.unpack("<I", msg2[8:12])[0]
    if msg_type != 2:
        raise Exception(f"msg_type not valid:{msg_type}")
    negotiate_flags = struct.unpack("<I", msg2[20:24])[0]
    server_challenge = msg2[24:32]
    return server_challenge, negotiate_flags


def create_ntlm_authenticate_message(server_challenge, user, domain, password, negotiate_flags):
    # 从配置信息取出数据
    is_unicode = negotiate_flags & NTLMSSP_NEGOTIATE_UNICODE
    is_negotiate_extended_security = negotiate_flags & NTLMSSP_NEGOTIATE_EXTENDED_SESSIONSECURITY
    flags = struct.pack('<I', NTLM_TYPE2_FLAGS)
    # 用于offset计算
    bytes_before_payload = 72
    payload_cur_offset = bytes_before_payload  # in bytes
    # 生成工作区、域名、用户名字节数组
    # 没设置NTLMSSP_NEGOTIATE_KEY_EXCH, 故encrypted_random_session_key是全0bit
    if is_unicode:
        workstation = bytes(gethostname().upper(), 'utf-16-le')
        domain_name = bytes(domain.upper(), 'utf-16-le')
        user_name = bytes(user, 'utf-16-le')
        encrypted_random_session_key = bytes("", 'utf-16-le')
    else:
        workstation = bytes(gethostname().upper(), 'ascii')
        domain_name = bytes(domain.upper(), 'ascii')
        user_name = bytes(user, 'ascii')
        encrypted_random_session_key = b""

    if is_negotiate_extended_security:
        client_challenge = generate_client_challenge()
        lm_challenge_response = client_challenge + func_z_bit(16)
        nt_challenge_response = func_desl(create_nt_hashed_password_v1(password),
                                          hashlib.md5(server_challenge + client_challenge).digest()[0:8])
    else:
        lm_challenge_response = func_desl(create_lm_hashed_password_v1(password), server_challenge)
        nt_challenge_response = func_desl(create_nt_hashed_password_v1(password), server_challenge)

    domain_name_len = struct.pack('<H', len(domain_name))
    domain_name_offset = struct.pack('<I', payload_cur_offset)
    payload_cur_offset += len(domain_name)

    user_name_len = struct.pack('<H', len(user_name))
    user_name_offset = struct.pack('<I', payload_cur_offset)
    payload_cur_offset += len(user_name)

    workstation_len = struct.pack('<H', len(workstation))
    workstation_offset = struct.pack('<I', payload_cur_offset)
    payload_cur_offset += len(workstation)

    lm_challenge_response_len = struct.pack('<H', len(lm_challenge_response))
    lm_challenge_response_offset = struct.pack('<I', payload_cur_offset)
    payload_cur_offset += len(lm_challenge_response)

    nt_challenge_response_len = struct.pack('<H', len(nt_challenge_response))
    nt_challenge_response_offset = struct.pack('<I', payload_cur_offset)
    payload_cur_offset += len(nt_challenge_response)

    encrypted_random_session_key_len = struct.pack('<H', len(encrypted_random_session_key))
    encrypted_random_session_key_offset = struct.pack('<I', payload_cur_offset)
    payload_cur_offset += len(encrypted_random_session_key)
    negotiate_flags = flags

    return base64.b64encode(NTLMSSP_SIGNATURE + NTLM_AUTHENTICATE_MESSAGE_TYPE +
                            lm_challenge_response_len * 2 + lm_challenge_response_offset +
                            nt_challenge_response_len * 2 + nt_challenge_response_offset +
                            domain_name_len * 2 + domain_name_offset + user_name_len * 2 + user_name_offset +
                            workstation_len * 2 + workstation_offset +
                            encrypted_random_session_key_len * 2 + encrypted_random_session_key_offset +
                            negotiate_flags + get_version_str() + domain_name + user_name + workstation +
                            lm_challenge_response + nt_challenge_response + encrypted_random_session_key).decode()


def generate_client_challenge():
    """
    生成客户端随机数, 8位
    :return: 8位随机数
    """
    client_challenge = b""
    for _ in range(8):
        client_challenge += bytes((random.getrandbits(8),))
    return client_challenge


def func_desl(password_hash, server_challenge):
    """
    把16位的pas.sword_hash和8位的server_challenge，用DESL算出24位的加密结果
    参考[MS-NLMP] 第6章 DESL(K, D)
    :param password_hash: password_hash
    :param server_challenge: server_challenge
    :return: 24位的加密结果
    """
    # password_hash 补齐到21位，后面补0
    password_hash = password_hash + func_z_bit((21 - len(password_hash)))
    return func_des(password_hash[0:7], server_challenge[0:8]) + \
           func_des(password_hash[7:14], server_challenge[0:8]) + \
           func_des(password_hash[14:21], server_challenge[0:8])


def func_des(para_k, para_d):
    """
    把7位para_k和8位的para_d，DES加密成8位的结果
    :param para_k: para_k
    :param para_d: para_d
    :return: DES加密成8位的结果
    """
    from Crypto.Cipher import DES

    cipher = DES.new(key56_to_key64(para_k), DES.MODE_ECB)
    ciphertext = cipher.encrypt(para_d)
    return ciphertext


def key56_to_key64(key):
    """
    用于DES加密的KEY，7字节转8字节
    :param key: 7字节key
    :return: 8字节key
    """
    if len(key) != 7:
        raise ValueError("DES 7-byte key is not 7 bytes in length, actual: %d" % len(key))

    new_key = b""
    for i in range(0, 8):
        if i == 0:
            new_value = struct.unpack("B", key[i: i + 1])[0]
        elif i == 7:
            new_value = struct.unpack("B", key[6:7])[0]
            new_value = (new_value << 1) & 0xFF
        else:
            new_value = struct.unpack("B", key[i - 1: i])[0]
            next_value = struct.unpack("B", key[i: i + 1])[0]
            new_value = ((new_value << (8 - i)) & 0xFF) | next_value >> i
        new_value = new_value & ~(1 << 0)
        new_value = new_value | int(not bit_count(new_value) & 0x1)
        new_key += struct.pack("B", new_value)

    return new_key


def bit_count(i):
    """
    计算比特1的数量
    :param i: 数字
    :return: 数量
    """
    count = 0
    while i:
        i &= i - 1
        count += 1

    return count


def func_z_bit(byte_length):
    """
    生成全0字节
    :param byte_length: 字节长度
    :return: 全0字节
    """
    return b'\0' * byte_length


def func_z_str(str_length):
    """
    生成全0字符
    :param str_length: 字符长度
    :return: 全0字节
    """
    return '\0' * str_length


def create_lm_hashed_password_v1(passwd_upper):
    """
    生成LM V1加密字符串
    :param passwd_upper: 密码
    :return: LM V1加密字符串
    """
    # 如果密码是hash格式，用:分割开，则取前半段
    if re.match(r'^[\w]{32}:[\w]{32}$', passwd_upper):
        return binascii.unhexlify(passwd_upper.split(':')[0])
    # 密码调整到14字符，超出忽略，不足后面补\b00
    passwd_upper = passwd_upper.upper()
    passwd_upper = passwd_upper + func_z_str(14 - len(passwd_upper))
    passwd_14bits = bytes(passwd_upper[0:14], 'utf8')
    return func_des(passwd_14bits[0:7], MAGIC_STR) + func_des(passwd_14bits[7:14], MAGIC_STR)


def create_nt_hashed_password_v1(passwd):
    """
    生成NT V1加密字符串
    :param passwd: 密码
    :return: NT V1加密字符串
    """
    # 如果密码是hash格式，用:分割开，则取前半段
    if re.match(r'^[\w]{32}:[\w]{32}$', passwd):
        return binascii.unhexlify(passwd.split(':')[1])
    return hashlib.new('md4', passwd.encode('utf-16le')).digest()
